{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from torch import multiprocessing\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torch.optim as optim\n",
    "from torcheval.metrics import MulticlassAccuracy\n",
    "\n",
    "is_fork = multiprocessing.get_start_method() == \"fork\"\n",
    "device = (\n",
    "    torch.device(0)\n",
    "    if torch.cuda.is_available() and not is_fork\n",
    "    else torch.device(\"cpu\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "transform = transforms.Compose(\n",
    "    [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]\n",
    ")\n",
    "\n",
    "batch_size = 32\n",
    "\n",
    "trainset = torchvision.datasets.CIFAR10(\n",
    "    root=\"./data\", train=True, download=True, transform=transform\n",
    ")\n",
    "trainloader = torch.utils.data.DataLoader(\n",
    "    trainset, batch_size=batch_size, shuffle=True, num_workers=2, drop_last=True\n",
    ")\n",
    "\n",
    "testset = torchvision.datasets.CIFAR10(\n",
    "    root=\"./data\", train=False, download=True, transform=transform\n",
    ")\n",
    "testloader = torch.utils.data.DataLoader(\n",
    "    testset, batch_size=batch_size, shuffle=False, num_workers=2\n",
    ")\n",
    "\n",
    "classes = (\n",
    "    \"plane\",\n",
    "    \"car\",\n",
    "    \"bird\",\n",
    "    \"cat\",\n",
    "    \"deer\",\n",
    "    \"dog\",\n",
    "    \"frog\",\n",
    "    \"horse\",\n",
    "    \"ship\",\n",
    "    \"truck\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Isomorphic_model(torch.nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        processing_channel=8,\n",
    "        horizontal_channel=8,\n",
    "        vertical_channel=8,\n",
    "        output_channel=16,\n",
    "        representation_channel=8,\n",
    "        kernel_size_1=3,\n",
    "        kernel_size_2=13,\n",
    "    ):\n",
    "        super(Isomorphic_model, self).__init__()\n",
    "        self.processing = torch.nn.LazyConv2d(\n",
    "            processing_channel,\n",
    "            kernel_size_1,\n",
    "            padding=\"same\",\n",
    "        )\n",
    "        self.horizontal_processing = torch.nn.LazyConv2d(\n",
    "            horizontal_channel,\n",
    "            (\n",
    "                kernel_size_2,\n",
    "                kernel_size_1,\n",
    "            ),\n",
    "            padding=\"same\",\n",
    "        )\n",
    "        self.vertical_processing = torch.nn.LazyConv2d(\n",
    "            vertical_channel,\n",
    "            (\n",
    "                kernel_size_1,\n",
    "                kernel_size_2,\n",
    "            ),\n",
    "            padding=\"same\",\n",
    "        )\n",
    "        self.output_processing = torch.nn.LazyConv2d(\n",
    "            output_channel, kernel_size_1, padding=\"same\"\n",
    "        )\n",
    "        self.representation_processing = torch.nn.LazyConv2d(\n",
    "            representation_channel, kernel_size_1, padding=\"same\"\n",
    "        )\n",
    "        self.bn1 = torch.nn.LazyBatchNorm2()\n",
    "        self.bn2 = torch.nn.LazyBatchNorm2()\n",
    "\n",
    "    def forward(self, x):\n",
    "        layer_4_output = torch.relu(self.processing(self.bn1(x)))\n",
    "        layer_2_3_output = torch.relu(\n",
    "            self.bn2(\n",
    "                torch.concat(\n",
    "                    [\n",
    "                        self.horizontal_processing(layer_4_output),\n",
    "                        self.vertical_processing(layer_4_output),\n",
    "                    ],\n",
    "                    axis=1,\n",
    "                )\n",
    "            )\n",
    "        )\n",
    "        layer_5_output = torch.sigmoid(self.output_processing(layer_2_3_output))\n",
    "        layer_6_output = torch.sigmoid(self.representation_processing(layer_2_3_output))\n",
    "        return layer_5_output, layer_6_output\n",
    "\n",
    "\n",
    "class MyModel(torch.nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        processing_channel=8,\n",
    "        horizontal_channel=8,\n",
    "        vertical_channel=8,\n",
    "        output_channel=16,\n",
    "        representation_channel=8,\n",
    "        kernel_size_1=3,\n",
    "        kernel_size_2=13,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        self.layers = torch.nn.ModuleList(\n",
    "            [\n",
    "                Isomorphic_model(\n",
    "                    processing_channel=processing_channel * 2**i,\n",
    "                    horizontal_channel=horizontal_channel * 2**i,\n",
    "                    vertical_channel=vertical_channel * 2**i,\n",
    "                    output_channel=output_channel * 2**i,\n",
    "                    representation_channel=representation_channel,\n",
    "                    kernel_size_1=kernel_size_1,\n",
    "                    kernel_size_2=kernel_size_2,\n",
    "                )\n",
    "                for i in range(1, 4)\n",
    "            ]\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        y = x * 1\n",
    "        for i in self.layers:\n",
    "            y, r = i(y)\n",
    "        return y, r\n",
    "\n",
    "\n",
    "model = MyModel().to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Isomorphic_model(torch.nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        processing_channel=8,\n",
    "        horizontal_channel=8,\n",
    "        vertical_channel=8,\n",
    "        output_channel=16,\n",
    "        representation_channel=8,\n",
    "        kernel_size_1=3,\n",
    "        kernel_size_2=13,\n",
    "    ):\n",
    "        super(Isomorphic_model, self).__init__()\n",
    "        self.fn = torch.nn.LazyLinear(10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        layer_6_output = torch.sigmoid(self.fn(x.view(-1, 32 * 32 * 3)))\n",
    "        return layer_6_output\n",
    "\n",
    "\n",
    "model = Isomorphic_model().to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.0001)\n",
    "loss = torch.nn.BCELoss()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,     2] loss: 0.000\n",
      "tensor(0.1163)\n",
      "[2,     2] loss: 0.000\n",
      "tensor(0.2998)\n",
      "[3,     2] loss: 0.000\n",
      "tensor(0.3602)\n",
      "[4,     2] loss: 0.000\n",
      "tensor(0.3972)\n",
      "[5,     2] loss: 0.000\n",
      "tensor(0.4335)\n",
      "[6,     2] loss: 0.000\n",
      "tensor(0.4573)\n",
      "[7,     2] loss: 0.000\n",
      "tensor(0.4756)\n",
      "[8,     2] loss: 0.000\n",
      "tensor(0.4906)\n",
      "[9,     2] loss: 0.000\n",
      "tensor(0.5028)\n",
      "[10,     2] loss: 0.000\n",
      "tensor(0.5161)\n",
      "[11,     2] loss: 0.000\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(100):\n",
    "    running_loss = 0.0\n",
    "    metric = MulticlassAccuracy()\n",
    "    for i, data in enumerate(trainloader, 0):\n",
    "        inputs, labels = data\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(inputs.to(device))[0][:, :10, 18, 18]\n",
    "        loss = torch.mean(\n",
    "            (torch.nn.functional.one_hot(labels, num_classes=10).to(device) - outputs)\n",
    "            ** 2\n",
    "        )\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        running_loss += loss.item()\n",
    "        metric.update(torch.argmax(outputs, dim=-1), labels)\n",
    "        if i % 11000 == 1:\n",
    "            print(f\"[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}\")\n",
    "            running_loss = 0.0\n",
    "    print(metric.compute())\n",
    "print(\"Finished Training\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1,     1] loss: 0.000\n",
      "[1,  1001] loss: 0.010\n",
      "tensor(0.1008)\n",
      "[2,     1] loss: 0.000\n",
      "[2,  1001] loss: 0.009\n",
      "tensor(0.1005)\n",
      "Finished Training\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(2):\n",
    "    running_loss = 0.0\n",
    "    metric = MulticlassAccuracy()\n",
    "    for i, data in enumerate(trainloader, 0):\n",
    "        inputs, labels = data\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(inputs.to(device))\n",
    "        loss = (\n",
    "            torch.mean(\n",
    "                torch.nn.functional.one_hot(labels, num_classes=10).to(device) - outputs\n",
    "            )\n",
    "            ** 2\n",
    "        )\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        running_loss += loss.item()\n",
    "        metric.update(torch.argmax(outputs, dim=-1), labels)\n",
    "        if i % 1000 == 0:\n",
    "            print(f\"[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}\")\n",
    "            running_loss = 0.0\n",
    "    print(metric.compute())\n",
    "print(\"Finished Training\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
